﻿#include "pch.h"
#include "FloppyDisk.h"

FloppyDisk::FloppyDisk(UINT nTracks, UINT nSectorPerTrack, UINT nSectorSize)
    : m_nTracks(nTracks)
    , m_nSectorPerTrack(nSectorPerTrack)
    , m_nSectorSize(nSectorSize)
{
    m_nSectorCount = nTracks * nSectorPerTrack;
    m_nCapacity = m_nSectorCount * nSectorSize;
}

FloppyDisk::~FloppyDisk()
{}

UINT FloppyDisk::GetTrackCount() const
{
    return m_nTracks;
}

UINT FloppyDisk::GetSectorCount() const
{
    return m_nSectorCount;
}

UINT FloppyDisk::GetSectorPerTrack() const
{
    return m_nSectorPerTrack;
}

UINT FloppyDisk::GetSectorSize() const
{
    return m_nSectorSize;
}

UINT FloppyDisk::GetCapacity() const
{
    return m_nCapacity;
}

BOOL FloppyDisk::Seek(UINT nTrack, UINT nSector)
{
    if(nTrack < m_nTracks && nSector >= 1 && nSector <= m_nSectorPerTrack) {
        m_nCurrentTrackNumber = nTrack;
        m_nCurrentSectorNumber = nSector;
        return TRUE;
    }
    return FALSE;
}

BOOL FloppyDisk::WriteImage(
    const CString& pszDevicePath,
    const CString& pszFilePath,
    PUINT pNumberOfTracksAffected,
    PUINT pNumberOfSectorsAffected,
    PUINT pNumberOfBytesWritten,
    CString& pszError)
{
    BOOL bRet = TRUE;
    HANDLE hDevice = INVALID_HANDLE_VALUE;
    HANDLE hFile = INVALID_HANDLE_VALUE;
    LPVOID lpBuffer = NULL;

    do {
        // Open device
        hDevice = CreateFile(
            (LPCTSTR)pszDevicePath,
            GENERIC_WRITE,
            FILE_SHARE_READ | FILE_SHARE_WRITE,
            NULL,
            OPEN_EXISTING,
            0,
            NULL
        );

        if(hDevice == INVALID_HANDLE_VALUE) {
            bRet = FALSE;
            pszError.Format(_T("Fail to open device: %s"), (LPCTSTR)pszDevicePath);
            break;
        }

        // Open file
        hFile = CreateFile(
            pszFilePath,
            GENERIC_READ,
            FILE_SHARE_READ,
            NULL,
            OPEN_EXISTING,
            FILE_ATTRIBUTE_NORMAL,
            NULL
        );

        if(hFile == INVALID_HANDLE_VALUE) {
            bRet = FALSE;
            pszError.Format(_T("Fail to open image file: %s"), (LPCTSTR)pszFilePath);
            break;
        }

        // If file size > 1M, not supported
        if(GetFileSize(hFile, NULL) > 1024 * 1024) {
            bRet = FALSE;
            pszError.Format(_T("File size must be within: 1024*1024 bytes"));
            break;
        }

        (*pNumberOfBytesWritten) = 0;
        (*pNumberOfTracksAffected) = 0;
        (*pNumberOfSectorsAffected) = 0;
        UINT nTrackNumber = m_nCurrentTrackNumber;
        UINT nSectorNumber = m_nCurrentSectorNumber;
        UINT FLASH_SECTOR_SIZE = 512;
        DWORD dwBytesRead = 0;

        // Sector write buffer
        lpBuffer = HeapAlloc(GetProcessHeap(), 0, FLASH_SECTOR_SIZE);
        ZeroMemory(lpBuffer, FLASH_SECTOR_SIZE);

        // Read file
        while(ReadFile(hFile, lpBuffer, m_nSectorSize, &dwBytesRead, NULL) && dwBytesRead > 0) {
            // Next track
            if(nSectorNumber > m_nSectorPerTrack) {
                nSectorNumber = 1;
                nTrackNumber++;
                (*pNumberOfTracksAffected)++;
            }

            // Write block of 512 bytes
            UINT nBlockNumber = GetLogicalBlockAddress(nTrackNumber, nSectorNumber);
            UINT nNumberOfBytesWritten = 0;
            if(!WriteBlock(hDevice, nBlockNumber, FLASH_SECTOR_SIZE, lpBuffer, &nNumberOfBytesWritten)) {
                bRet = FALSE;
                pszError.Format(_T("Fail to write data, block number: %d"), nBlockNumber);
                break;
            }

            // Next sector
            nSectorNumber++;
            (*pNumberOfBytesWritten) += dwBytesRead;
            (*pNumberOfSectorsAffected)++;
            ZeroMemory(lpBuffer, FLASH_SECTOR_SIZE);
        }

    } while(0);

    if(hDevice != INVALID_HANDLE_VALUE)
        CloseHandle(hDevice);

    if(hFile != INVALID_HANDLE_VALUE)
        CloseHandle(hFile);

    if(lpBuffer != NULL)
        HeapFree(GetProcessHeap(), 0, lpBuffer);

    if(bRet && (*pNumberOfTracksAffected) == 0 && (*pNumberOfSectorsAffected > 0))
        (*pNumberOfTracksAffected) = 1;

    return bRet;
}

BOOL FloppyDisk::WriteMBR(const CString& pszDevicePath,
    const CString& pszFilePath,
    PUINT pNumberOfBytesWritten,
    CString& pszError)
{
    BOOL bRet = TRUE;
    HANDLE hDevice = INVALID_HANDLE_VALUE;
    HANDLE hFile = INVALID_HANDLE_VALUE;
    LPVOID lpBuffer = NULL;

    do {
        hDevice = CreateFile(
            (LPCTSTR)pszDevicePath,
            GENERIC_READ | GENERIC_WRITE,
            FILE_SHARE_READ | FILE_SHARE_WRITE,
            NULL,
            OPEN_EXISTING,
            0,
            NULL
        );

        if(hDevice == INVALID_HANDLE_VALUE) {
            bRet = FALSE;
            pszError.Format(_T("Fail to open device: %s"), (LPCTSTR)pszDevicePath);
            break;
        }

        hFile = CreateFile(
            pszFilePath,
            GENERIC_READ,
            FILE_SHARE_READ,
            NULL,
            OPEN_EXISTING,
            FILE_ATTRIBUTE_NORMAL,
            NULL
        );

        if(hFile == INVALID_HANDLE_VALUE) {
            bRet = FALSE;
            pszError.Format(_T("Fail to open image file: %s"), (LPCTSTR)pszFilePath);
            break;
        }

        // If file size > 512 bytes, not supported
        if(GetFileSize(hFile, NULL) > 512) {
            bRet = FALSE;
            pszError.Format(_T("Size of bootloader file must be within: 512 bytes"));
            break;
        }

        (*pNumberOfBytesWritten) = 0;
        UINT nTrackNumber = m_nCurrentTrackNumber;
        UINT nSectorNumber = m_nCurrentSectorNumber;
        UINT FLASH_SECTOR_SIZE = 512;
        DWORD dwBytesRead = 0;

        lpBuffer = HeapAlloc(GetProcessHeap(), 0, FLASH_SECTOR_SIZE);
        ZeroMemory(lpBuffer, FLASH_SECTOR_SIZE);

        // Read file
        if(ReadFile(hFile, lpBuffer, FLASH_SECTOR_SIZE, &dwBytesRead, NULL) && dwBytesRead > 0) {

            UINT nNumberOfBytesWritten = 0;
            if(!WriteBlock(hDevice, 0, FLASH_SECTOR_SIZE, lpBuffer, &nNumberOfBytesWritten)) {
                bRet = FALSE;
                pszError.Format(_T("Fail to write data to block: 0"));
                break;
            }
            (*pNumberOfBytesWritten) += dwBytesRead;
        }

    } while(0);

    if(hDevice != INVALID_HANDLE_VALUE)
        CloseHandle(hDevice);

    if(hFile != INVALID_HANDLE_VALUE)
        CloseHandle(hFile);

    if(lpBuffer != NULL)
        HeapFree(GetProcessHeap(), 0, lpBuffer);

    return bRet;
}

UINT FloppyDisk::GetLogicalBlockAddress(UINT nTrackNumber, UINT nSectorNumber) const
{
    UINT nLBA_High = 0;
    UINT nLBA_Mid = nTrackNumber;
    UINT nLBA_Low = nSectorNumber;
    return (nLBA_High << 16) & 0xFF0000 | (nLBA_Mid << 8) & 0xFF00 | (nLBA_Low & 0xFF);
}

BOOL FloppyDisk::WriteBlock(HANDLE hDevice, UINT nBlockNumber, UINT nBlockSize,
    LPVOID lpBuffer, PUINT pNumberOfBytesWritten)
{
    ASSERT(nBlockNumber >= 0 && nBlockSize > 0);
    if(hDevice == INVALID_HANDLE_VALUE)
        return FALSE;

    LARGE_INTEGER liDistanceToMove = { 0 };
    liDistanceToMove.QuadPart = nBlockNumber * nBlockSize;

    if(!SetFilePointerEx(hDevice, liDistanceToMove, NULL, FILE_BEGIN))
        return FALSE;

    DWORD bytesWritten;
    if(!WriteFile(hDevice, lpBuffer, nBlockSize, &bytesWritten, NULL)) {
        CString pszError;
        pszError.Format(_T("WriteFile Win32 Error: %d"), GetLastError());
        AfxMessageBox(pszError);
        return FALSE;
    }

    *pNumberOfBytesWritten = bytesWritten;
    return TRUE;
}